from typing import Any, Optional

from .solver_wrapper_base import SolverWrapperBase, ProblemDefinitionBase, FENICSX_AVAILABLE
from .simulation_parameters import LinearSolverSettings, SolverSettings # 确保导入 LinearSolverSettings

if FENICSX_AVAILABLE:
    from mpi4py import MPI
    from dolfinx import fem
    from dolfinx.fem.petsc import LinearProblem # 只导入 LinearProblem
    from petsc4py import PETSc # KSP 通常直接来自 petsc4py.PETSc，或者通过 LinearProblem 实例的 .ksp 属性访问
else:
    # 虚拟类型
    class MPI:
        class Comm: pass
    class fem:
        class Function: pass
    class LinearProblem: # 模拟 dolfinx.fem.petsc.LinearProblem
        def __init__(self, a, L, u, bcs=[], petsc_options={}): 
            self._ksp = KSP() # 模拟内部的 KSP
        def solve(self): 
            # 模拟求解，返回一个模拟的解对象 (或直接修改传入的 u)
            # return "mock_solution_function"
            self._ksp.iterations +=1 # 模拟迭代次数增加
            self._ksp.converged = True
            return fem.Function() # 返回一个模拟的 fem.Function 作为解
        @property
        def ksp(self):
            return self._ksp
            
    class KSP: # 模拟 PETSc KSP
        def __init__(self):
            self.iterations = 0
            self.converged = False
            self.rtol = 0.0
            self.atol = 0.0
            self.max_it = 0
            self.divtol = 0.0

        def setType(self, ksp_type): pass
        def getPC(self): return MockPC()
        def setTolerances(self, rtol=None, atol=None, divtol=None, max_it=None): 
            if rtol is not None: self.rtol = rtol
            if atol is not None: self.atol = atol
            if divtol is not None: self.divtol = divtol
            if max_it is not None: self.max_it = max_it
        def setOptionsPrefix(self, prefix): pass
        def appendOptionsPrefix(self, prefix): pass
        def setFromOptions(self): pass
        def getIterationNumber(self): return self.iterations
        def getConvergedReason(self): return 1 if self.converged else -1 # 模拟 PETSc 收敛原因
    
    class MockPC: # 模拟 PETSc PC
        def setType(self, pc_type): pass

    # class LinearSolverSettings: pass # 已从 simulation_parameters 导入

class LinearSolverWrapper(SolverWrapperBase):
    """
    封装 dolfinx.fem.petsc.LinearProblem (及其内部的 PETSc KSP) 的线性求解器包装器。
    """

    def __init__(self,
                 problem_definition: ProblemDefinitionBase,
                 solver_params: SolverSettings, # 使用通用的 SolverSettings
                 comm: Optional[MPI.Comm]):
        """
        初始化 LinearSolverWrapper。

        参数:
            problem_definition (ProblemDefinitionBase):
                此求解器将要解决的物理问题的定义实例。
            solver_params (SolverSettings):
                包含求解器配置的参数对象。线性特定设置应在其 .linear_solver_config 属性中。
            comm (Optional[MPI.Comm]):
                MPI 通信器。
        """
        super().__init__(problem_definition, solver_params, comm)
        self._is_initialized_correctly = False
        self.linear_solver_settings: Optional[LinearSolverSettings] = self.solver_params.linear_solver_config

        if self.linear_solver_settings is None:
            print(f"警告: {self.problem_definition.problem_name} 的线性求解器配置缺失。将使用默认或模拟行为。")
            # self.linear_solver_settings = LinearSolverSettings() # 如果需要，创建默认配置

        if not FENICSX_AVAILABLE:
            print("FEniCSx 不可用，将使用模拟的线性求解器对象进行 LinearSolverWrapper。")
            # _solver_object 在这种情况下是一个模拟的 LinearProblem
            self._solver_object = LinearProblem(None, None, fem.Function()) # 模拟构造
            self._is_initialized_correctly = True
            return

        # FEniCSx 可用时的逻辑
        if self._comm is None:
            print(f"警告: MPI 通信器未提供给 LinearSolverWrapper '{self.problem_definition.problem_name}'，求解器对象未创建。")
            return

        print(f"[LWrapper __init__] FENICSX_AVAILABLE (internal): {FENICSX_AVAILABLE}") # 新增打印

        a_form = self.problem_definition.get_bilinear_form()
        L_form = self.problem_definition.get_linear_form()
        u_solve = self.problem_definition.get_variable_to_solve() # LinearProblem 需要一个 fem.Function 来存储解
        
        print(f"[LWrapper __init__] Problem: {self.problem_definition.problem_name}, u_solve type from problem_def: {type(u_solve)}") # 新增打印
        if hasattr(u_solve, 'name'): print(f"[LWrapper __init__] u_solve name: {u_solve.name}") # 新增打印

        bcs = self.problem_definition.get_boundary_conditions()

        if a_form is None:
            print(f"错误: 双线性形式 a 未在问题定义 '{self.problem_definition.problem_name}' 中提供，无法创建线性求解器。")
            return
        if L_form is None:
            print(f"错误: 线性形式 L 未在问题定义 '{self.problem_definition.problem_name}' 中提供，无法创建线性求解器。")
            return
        if u_solve is None: # 虽然基类保证了 get_variable_to_solve 不会是 None，但以防万一
            print(f"错误: 求解变量 u_solve 未在问题定义 '{self.problem_definition.problem_name}' 中提供，无法创建线性求解器。")
            return

        # 从 linear_solver_settings 中获取自定义 PETSc 选项
        # 旧代码: petsc_options = self.linear_solver_settings.petsc_options if self.linear_solver_settings else {}
        petsc_options = self.linear_solver_settings.custom_options if self.linear_solver_settings and self.linear_solver_settings.custom_options is not None else {}

        try:
            # 确保 u_solve 是一个 fem.Function，LinearProblem 需要它来存储解
            if FENICSX_AVAILABLE and not isinstance(u_solve, fem.Function):
                print(f"错误: [LWrapper Check] 求解变量 'u_solve' ({type(u_solve).__name__}) 不是一个有效的 dolfinx.fem.Function 类型 (FENICSX_AVAILABLE=True)，无法创建线性求解器。") # 修改打印
                return
            # 当 FENICSX_AVAILABLE is False, u_solve 会是 MockFEMFunction，而 fem.Function 是模拟的。
            # 此时 isinstance 检查可能不适用或导致问题，所以仅在 FENICSX_AVAILABLE 时严格检查。
            
            # 调试信息
            print(f"[调试 LinearSolverWrapper] FENICSX_AVAILABLE: {FENICSX_AVAILABLE}")
            print(f"[调试 LinearSolverWrapper] a_form type: {type(a_form)}, value: {a_form}")
            print(f"[调试 LinearSolverWrapper] L_form type: {type(L_form)}, value: {L_form}")
            print(f"[调试 LinearSolverWrapper] u_solve type: {type(u_solve)}, name: {u_solve.name if hasattr(u_solve, 'name') else 'N/A'}")
            print(f"[调试 LinearSolverWrapper] bcs[0] type: {type(bcs[0]) if bcs else 'N/A'}") # Handle empty bcs
            print(f"[调试 LinearSolverWrapper] petsc_options: {petsc_options}")

            # 模仿 Dokken 教程: a, L 位置参数, bcs 和 petsc_options 关键字参数
            # 不在构造时传入 u_solve
            self._solver_object = LinearProblem(a_form, L_form, bcs=bcs, petsc_options=petsc_options)
            
            # 新增诊断：检查 LinearProblem 实例及其 A 和 b 属性
            print(f"[LWrapper Post-LP-Init] _solver_object (LinearProblem instance): {self._solver_object}")
            if hasattr(self._solver_object, 'A'):
                print(f"[LWrapper Post-LP-Init] _solver_object.A: {self._solver_object.A}")
                # print(f"[LWrapper Post-LP-Init] _solver_object.A type: {type(self._solver_object.A)}")
            else:
                print("[LWrapper Post-LP-Init] _solver_object has NO 'A' attribute.")

            if hasattr(self._solver_object, 'b'):
                print(f"[LWrapper Post-LP-Init] _solver_object.b: {self._solver_object.b}")
                # print(f"[LWrapper Post-LP-Init] _solver_object.b type: {type(self._solver_object.b)}")
            else:
                print("[LWrapper Post-LP-Init] _solver_object has NO 'b' attribute.")

            # 检查 _solver_object 是否成功创建并且具有 ksp 属性
            # 直接尝试访问 .ksp，如果属性不存在或获取时出错，会抛出 AttributeError
            try:
                self._ksp = self._solver_object.ksp 
            except AttributeError:
                print(f"[LWrapper KSP Access] .ksp 属性不存在，尝试访问 .solver 属性作为备选。")
                try:
                    self._ksp = self._solver_object.solver
                except AttributeError:
                    print(f"错误: LinearProblem 对象 '{self.problem_definition.problem_name}' 创建后既不具有 'ksp' 也不具有 'solver' 属性。")
                    self._is_initialized_correctly = False
                    return # 如果两个都不存在，则初始化失败
            
            if self._ksp is None:
                print(f"错误: LinearProblem 对象 '{self.problem_definition.problem_name}' 的 .ksp 属性为 None。")
                self._is_initialized_correctly = False
                return # 如果 ksp 为 None，则初始化失败
            
            self._ksp.setOptionsPrefix(f"{self.problem_definition.problem_name}_")
            # 应用来自 solver_params 的 KSP 设置
            if self.linear_solver_settings:
                self._ksp.setTolerances(rtol=self.linear_solver_settings.relative_tolerance,
                                        atol=self.linear_solver_settings.absolute_tolerance,
                                        divtol=self.linear_solver_settings.divergence_tolerance,
                                        max_it=self.linear_solver_settings.max_iterations)
                if self.linear_solver_settings.initial_guess_nonzero:
                    self._ksp.setInitialGuessNonzero(True)
                # KSP的自定义选项通常通过 options DB 在命令行或代码中设置，
                # LinearProblem 的 petsc_options 应该已经处理了传给KSP的选项。
                # 如果需要，可以考虑在这里用 self._ksp.setFromOptions()，但要确保选项前缀正确。
            self._is_initialized_correctly = True
            print(f"LinearSolverWrapper '{self.problem_definition.problem_name}' 初始化成功。KSP 对象已获取并配置。")

        except AttributeError as e:
            print(f"错误: 初始化线性求解器时发生 AttributeError: {e}。可能是 FEniCSx 版本不兼容或 PETSc 配置问题。")
        except Exception as e:
            print(f"错误: 初始化线性求解器时发生未知错误: {e}")

    def solve(self) -> bool:
        """
        使用配置的线性求解器 (dolfinx.fem.petsc.LinearProblem) 求解线性问题。

        返回:
            bool: 如果求解过程成功 (通常线性求解器执行一次)，则返回 True。
                  如果 KSP 迭代未能收敛 (虽然 LinearProblem.solve() 不直接返回此信息，
                  但可以通过 KSP 对象检查)，也可能返回 False。
        """
        if not self._is_initialized_correctly or self._solver_object is None:
            print(f"错误: 线性求解器 '{self.problem_definition.problem_name}' 未正确初始化或不可用，无法执行 solve()。")
            self._is_converged = False
            self._num_iterations = 0
            return False

        try:
            print(f"调用 dolfinx.fem.petsc.LinearProblem.solve() 处理问题 '{self.problem_definition.problem_name}'...")
            # LinearProblem.solve() 返回解函数
            # 我们需要将解函数的值赋给 problem_definition 中的变量
            solved_u_function = self._solver_object.solve()
            target_u_function = self.problem_definition.get_variable_to_solve()

            if target_u_function is not None and hasattr(target_u_function, 'x') and hasattr(solved_u_function, 'x'):
                # 确保类型兼容或进行适当的转换/插值
                if target_u_function.function_space == solved_u_function.function_space:
                    target_u_function.x.array[:] = solved_u_function.x.array[:]
                    target_u_function.x.scatter_forward() # 确保 ghost 值更新
                    print(f"  解已从 LinearProblem.solve() 返回的函数复制到 '{target_u_function.name if hasattr(target_u_function, 'name') else '未知变量'}'。")
                else:
                    print(f"错误: LinearProblem.solve() 返回的解函数空间与目标函数空间不匹配。需要插值。")
                    # 这里可以添加插值逻辑 target_u_function.interpolate(solved_u_function)
                    # 但首先确保基本流程正确
                    self._is_converged = False # 标记为未解决，因为解未正确传递
                    return False 
            else:
                print(f"错误: 无法将 LinearProblem.solve() 的解赋给目标变量 (可能为None或缺少.x属性)。")
                self._is_converged = False
                return False
            
            # LinearProblem.solve() 不直接返回收敛状态布尔值。
            # 我们需要检查 KSP 对象的收敛原因。
            # 确保使用 self._ksp (它可能是 .solver 或 .ksp)
            if self._ksp is not None:
                convergence_reason = self._ksp.getConvergedReason()
                print(f"  KSP 收敛原因代码: {convergence_reason}")
                self._is_converged = convergence_reason >= 0
                self._num_iterations = self._ksp.getIterationNumber()
            else:
                print(f"错误: KSP 对象 (self._ksp) 为 None，无法获取收敛状态。")
                self._is_converged = False # 无法确定收敛性
            
            if self._is_converged:
                print(f"  LinearProblem 求解完成。KSP 迭代次数: {self._num_iterations}。收敛原因: {convergence_reason}。解已更新到 '{self.problem_definition.get_variable_to_solve().name if hasattr(self.problem_definition.get_variable_to_solve(), 'name') else '未知变量'}' 中。")
            else:
                print(f"  LinearProblem 求解完成，但 KSP 未能收敛。KSP 迭代次数: {self._num_iterations}。收敛原因: {convergence_reason}。")
            
            # 确保 ProblemDefinition 中的变量被更新 (LinearProblem 的 u 参数应该就是它)
            # u_target = self.problem_definition.get_variable_to_solve()
            # u_target.x.array[:] = solution_function.x.array[:] # 通常 LinearProblem 会直接修改传入的 u

            return self._is_converged

        except AttributeError as e_attr:
            print(f"在 LinearSolverWrapper.solve() (LinearProblem) 期间发生属性错误: {e_attr}")
            import traceback
            traceback.print_exc()
            self._is_converged = False
        except Exception as e:
            print(f"在 LinearSolverWrapper.solve() (LinearProblem) 期间发生未知错误: {e}")
            self._is_converged = False
            return False

    def update_problem_forms(self) -> None:
        """
        对于 LinearProblem，如果其底层的矩阵 A 或向量 b 需要重建
        （例如，因为 a_form 或 L_form 依赖于已更改的 fem.Constant 或 fem.Function），
        则可能需要重新创建 LinearProblem 实例或其内部组件。
        然而，通常 UFL 形式会动态引用这些量，在 assemble 时获取最新值。
        如果 LinearProblem 缓存了矩阵，则可能需要某种形式的更新。
        目前的 dolfinx.fem.petsc.LinearProblem 似乎在每次 solve() 时重新组装（如果需要）。
        """
        super().update_problem_forms()
        if FENICSX_AVAILABLE and self._is_initialized_correctly:
            print(f"LinearSolverWrapper.update_problem_forms() 被调用为 '{self.problem_definition.problem_name}'。")
            # LinearProblem 内部的 KSP 的 setOperators 可能会被再次调用，
            # 或者 LinearProblem 对象本身被重建（如果形式的结构改变了）。
            # 简单的做法是，如果形式的 *定义* 改变了，就重建 LinearProblem。
            # 如果只是形式所依赖的 *值* 变了 (例如 fem.Constant)，LinearProblem 的 assemble 应该能处理。
            print("  对于 LinearProblem，通常不需要特殊操作，因为形式会在组装时更新。")
            # 如果确实需要重建，逻辑会像这样：
            # new_a = self.problem_definition.get_bilinear_form()
            # new_L = self.problem_definition.get_linear_form()
            # u_solve = self.problem_definition.get_variable_to_solve()
            # bcs = self.problem_definition.get_boundary_conditions()
            # self._solver_object = LinearProblem(new_a, new_L, u=u_solve, bcs=bcs, petsc_options=...)
            pass

if __name__ == '__main__':
    # 此 __main__ 部分仅用于基本演示
    print("LinearSolverWrapper 子类已定义。")

    class MockSimParamsLinear:
        class LinearConfig:
            solver_type = "gmres"
            preconditioner_type = "ilu"
            absolute_tolerance = 1e-10
            relative_tolerance = 1e-8
            max_iterations = 1000
            divergence_tolerance = 1e4
            monitor_convergence = False
            initial_guess_nonzero = False
            custom_options = None
        linear_solver_config = LinearConfig()

    class MockProblemDefLinear(ProblemDefinitionBase):
        def __init__(self, name, comm_inner):
            self.problem_name = name
            self.comm = comm_inner
            if FENICSX_AVAILABLE:
                # domain_mock = fem.create_unit_square(self.comm, 2, 2)
                # V_mock = fem.functionspace(domain_mock, ("CG", 1))
                # self._u_sol = fem.Function(V_mock)
                # u_trial_mock = ufl.TrialFunction(V_mock)
                # v_test_mock = ufl.TestFunction(V_mock)
                # self._a = ufl.inner(ufl.grad(u_trial_mock), ufl.grad(v_test_mock)) * ufl.dx
                # self._L = fem.Constant(domain_mock, 1.0) * v_test_mock * ufl.dx
                pass # 简化
            else:
                self._u_sol = fem.Function() # 模拟
                self._a = "mock_a_form"
                self._L = "mock_L_form"
            self._bcs = []

        def setup_problem(self): pass
        def get_bilinear_form(self): return self._a if hasattr(self, '_a') else "mock_a_form"
        def get_linear_form(self): return self._L if hasattr(self, '_L') else "mock_L_form"
        def get_boundary_conditions(self): return self._bcs
        def get_variable_to_solve(self): return self._u_sol if hasattr(self, '_u_sol') else fem.Function() # 模拟
        def get_test_function(self): return "mock_test_func"
        def get_trial_function(self): return "mock_trial_func"

    mock_mpi_comm_main_lin = MPI.COMM_WORLD if FENICSX_AVAILABLE else MPI.Comm()
    print(f"FEniCSx 可用: {FENICSX_AVAILABLE}")

    prob_def_lin_instance = MockProblemDefLinear("test_lin_problem", mock_mpi_comm_main_lin)
    solver_settings_lin_instance = MockSimParamsLinear()

    try:
        lin_solver = LinearSolverWrapper(prob_def_lin_instance, solver_settings_lin_instance, mock_mpi_comm_main_lin)
        print(f"创建的求解器: {lin_solver}")
        if lin_solver._is_initialized_correctly:
            lin_solver.solve()
            print(f"求解后: {lin_solver}")
        else:
            print("求解器未正确初始化，跳过求解。")
    except Exception as e_main_lin:
        print(f"运行 LinearSolverWrapper 示例时发生主错误: {e_main_lin}")
        import traceback
        traceback.print_exc() 